import numpy as np
import matplotlib.pyplot as plt
import scipy.integrate as inte

def sgn(x):
    xn=np.linalg.norm(x)
    if xn==0:
        return np.zeros_like(x)
    else:
        return x/xn
    
def rhs(t,x):
    d=np.array([np.sin(t),np.cos(t)])
    d2=np.array([1,1])
    return -2*sgn(x)+d2

class Sol:
    def __init__(self, t, y):
        self.t = t
        self.y = y


def solve_rk4(fun, tspan, y0, t_eval=None):
    # 解包时间区间
    t0, tf = tspan
    # 如果没有指定评估时间点，则默认在时间区间内均匀取点
    if t_eval is None:
        t_eval = np.linspace(t0, tf, 100)
    # 初始化解数组
    y = np.zeros((len(t_eval), len(y0)))
    y[0] = y0
    # 开始循环求解
    for i in range(len(t_eval) - 1):
        t = t_eval[i]
        h = t_eval[i + 1] - t
        k1 = fun(t, y[i])
        k2 = fun(t + h / 2, y[i] + h / 2 * k1)
        k3 = fun(t + h / 2, y[i] + h / 2 * k2)
        k4 = fun(t + h, y[i] + h * k3)
        y[i + 1] = y[i] + h / 6 * (k1 + 2 * k2 + 2 * k3 + k4)
    return Sol(t_eval, y.T)

sol=solve_rk4(rhs,[0,60],[1,0],t_eval=np.arange(0,60,0.0001))

#%%
# 绘制曲线
plt.figure(figsize=(10, 6))
plt.subplot(2,1,1)
plt.plot(sol.t, sol.y[1,:], label=r'$x$')
plt.grid(True)
plt.legend()
plt.subplot(2,1,2)
t=sol.t[-100:-1]
x=sol.y[:,-100:-1]
phis=[]
for it,ti in enumerate(t):
    xi=x[:,it]
    ui=-2*sgn(xi)
    phi=np.arctan2(xi[1],xi[0])
    phis.append(phi)
plt.plot(t,phis)
plt.show()
    
# %%
